#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Coinank WebSocket客户端
连接到coinank WebSocket服务，接收实时爆仓数据
"""

import asyncio
import json
import time
import websockets
from typing import Optional, Callable
from loguru import logger

class CoinankWebSocketClient:
    """Coinank WebSocket客户端"""
    
    def __init__(self, data_processor):
        self.url = "wss://ws.coinank.com/wsKline/wsKline"
        self.data_processor = data_processor
        self.websocket: Optional[websockets.WebSocketServerProtocol] = None
        self.is_running = False
        self.reconnect_interval = 5  # 重连间隔(秒)
        self.max_reconnect_attempts = 10
        self.reconnect_count = 0
        self.heartbeat_task = None
        self.heartbeat_interval = 30  # 30秒心跳间隔
        self.last_pong_time = None
        self.ping_timeout = 10  # ping超时时间
        
    async def start(self):
        """启动WebSocket客户端"""
        self.is_running = True
        while self.is_running:
            try:
                await self._connect_and_listen()
            except Exception as e:
                logger.error(f"WebSocket连接异常: {e}")
                if self.is_running:
                    await self._handle_reconnect()
                    
    async def _start_heartbeat(self):
        """启动心跳任务"""
        if self.heartbeat_task is None:
            self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())
            
    async def _stop_heartbeat(self):
        """停止心跳任务"""
        if self.heartbeat_task:
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass
            self.heartbeat_task = None
            
    async def _heartbeat_loop(self):
        """心跳循环"""
        logger.info(f"心跳循环已启动，间隔: {self.heartbeat_interval}秒")
        
        while self.is_running and self.is_connected():
            try:
                await asyncio.sleep(self.heartbeat_interval)
                
                if not self.is_connected():
                    logger.warning("WebSocket连接已断开，停止心跳")
                    break
                    
                # 发送ping
                ping_time = time.time()
                await self.websocket.send('ping')
                logger.debug(f"发送心跳ping，时间: {ping_time}")
                
                # 检查上次pong响应时间
                if self.last_pong_time and (ping_time - self.last_pong_time) > self.ping_timeout * 2:
                    logger.warning(f"心跳响应超时，上次pong时间: {self.last_pong_time}，当前时间: {ping_time}")
                    # 连接可能有问题，让主循环处理重连
                    break
                    
            except asyncio.CancelledError:
                logger.info("心跳循环被取消")
                break
            except Exception as e:
                logger.error(f"心跳发送失败: {e}")
                break
                
        logger.info("心跳循环已停止")
                    
    async def stop(self):
        """停止WebSocket客户端"""
        logger.info("正在停止WebSocket客户端...")
        self.is_running = False
        await self._stop_heartbeat()
        if self.websocket:
            await self.websocket.close()
            logger.info("WebSocket连接已关闭")
            
    def is_connected(self) -> bool:
        """检查WebSocket是否连接"""
        return self.websocket is not None and self.websocket.state.name == 'OPEN'
        
    async def _connect_and_listen(self):
        """连接并监听数据"""
        logger.info(f"正在连接到 {self.url}...")
        
        async with websockets.connect(
            self.url,
            ping_interval=20,
            ping_timeout=10,
            close_timeout=10
        ) as websocket:
            self.websocket = websocket
            self.reconnect_count = 0
            logger.info("WebSocket连接成功")
            
            # 启动心跳任务
            await self._start_heartbeat()
            
            # 发送订阅消息
            subscribe_msg = {
                "op": "subscribe",
                "args": "liqOrder@All@All@1m"
            }
            await websocket.send(json.dumps(subscribe_msg))
            logger.info(f"已发送订阅消息: {subscribe_msg}")
            
            # 监听消息
            async for message in websocket:
                try:
                    await self._handle_message(message)
                except Exception as e:
                    logger.error(f"处理消息时出错: {e}")
                    
    async def _handle_message(self, message: str):
        """处理接收到的消息"""
        try:
            # 处理ping消息
            if message == 'ping':
                await self.websocket.send('pong')
                logger.debug("收到服务器ping，已回复pong")
                return
                
            # 处理pong消息
            if message == 'pong':
                self.last_pong_time = time.time()
                logger.debug(f"收到服务器pong响应，时间: {self.last_pong_time}")
                return
            
            data = json.loads(message)
            
            # 检查消息类型
            if data.get("op") == "push" and data.get("success"):
                # 处理爆仓数据
                liquidation_data = data.get("data", [])
                if liquidation_data:
                    logger.debug(f"收到 {len(liquidation_data)} 条爆仓数据")
                    await self.data_processor.process_liquidation_data(liquidation_data)
            elif data.get("op") == "subscribe":
                # 订阅确认
                if data.get("success"):
                    logger.info(f"订阅成功: {data.get('args')}")
                else:
                    logger.error(f"订阅失败: {data}")
            else:
                logger.debug(f"收到其他消息: {data}")
                
        except json.JSONDecodeError as e:
            # 如果不是JSON格式，可能是纯文本ping/pong消息
            if message.strip() in ['ping', 'pong']:
                if message.strip() == 'ping':
                    await self.websocket.send('pong')
                    logger.debug("收到纯文本ping，已回复pong")
                elif message.strip() == 'pong':
                    self.last_pong_time = time.time()
                    logger.debug(f"收到纯文本pong响应，时间: {self.last_pong_time}")
            else:
                logger.error(f"JSON解析错误: {e}, 原始消息: {message[:200]}...")
        except Exception as e:
            logger.error(f"处理消息异常: {e}")
            
    async def _handle_reconnect(self):
        """处理重连逻辑"""
        if self.reconnect_count >= self.max_reconnect_attempts:
            logger.error(f"达到最大重连次数 {self.max_reconnect_attempts}，停止重连")
            self.is_running = False
            return
            
        self.reconnect_count += 1
        wait_time = min(self.reconnect_interval * self.reconnect_count, 60)
        logger.info(f"第 {self.reconnect_count} 次重连，等待 {wait_time} 秒...")
        await asyncio.sleep(wait_time)
        
    async def send_ping(self):
        """发送心跳包"""
        if self.is_connected():
            try:
                ping_msg = {"op": "ping"}
                await self.websocket.send(json.dumps(ping_msg))
                logger.debug("已发送心跳包")
            except Exception as e:
                logger.error(f"发送心跳包失败: {e}")
                
    async def get_connection_stats(self) -> dict:
        """获取连接统计信息"""
        return {
            "connected": self.is_connected(),
            "reconnect_count": self.reconnect_count,
            "is_running": self.is_running,
            "url": self.url
        }